from langchain import OpenAI, SQLDatabase, SQLDatabaseChain
from langchain.prompts.prompt import PromptTemplate
db = SQLDatabase.from_uri(
    "mysql+pymysql://root:password@115.159.119.116:3306/openai_demo",
    include_tables=['data_wide_mdm_myuser'],
    sample_rows_in_table_info=2
    )
llm = OpenAI(temperature=0)

_DEFAULT_TEMPLATE = """Given an input question, first create a syntactically correct {dialect} query to run, then look at the results of the query and return the answer.
Use the following format:

Question: "Question here"
SQLQuery: "SQL Query to run"
SQLResult: "Result of the SQLQuery"
Answer: "Final answer here"

Only use the following tables:

{table_info}

如果用户查询水果信息，其实代表用户表。

Question: {input}"""
PROMPT = PromptTemplate(
    input_variables=["input", "table_info", "dialect"], template=_DEFAULT_TEMPLATE
)
db_chain = SQLDatabaseChain.from_llm(llm, db, prompt=PROMPT, verbose=True)
db_chain.run("有多少用户?")
db_chain.run("有多少水果?")